import numpy as np
import argparse
import cv2
import os
import torch
from  onnx_infer import infer

def image_recognize(img_path, save_path):
    # Find haar cascade to draw bounding box around face
    frame = cv2.imread(img_path)
    #模型推理
    result = infer(frame)
    now_path = os.getcwd()

    cv2.imwrite(save_path,result)

    return save_path





#屌用摄像头函数
def video_recognize(v_path, s_path):
    cap = cv2.VideoCapture(v_path)
    fps = cap.get(cv2.CAP_PROP_FPS)  # 帧率
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 宽
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 高
    fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')  # 指定视频编码方式
    videoWriter = cv2.VideoWriter(s_path, fourcc, fps, (w, h))  # 创建视频写对象
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # 视频总帧数

    if v_path == 0:
        while 1:
            # Find haar cascade to draw bounding box around face
            ret, frame = cap.read()
            #frame = cv2.flip(frame, 1)
            #result = get_results(model, device, fire_img, frame.copy())


            # show the output frame
            cv2.imshow("Frame", frame)
            key = cv2.waitKey(1) & 0xFF

            # if the `q` key was pressed, break from the loop
            if key == ord("q"):
                break
    else:
        while 1:
            # Find haar cascade to draw bounding box around face
            ret, frame = cap.read()
            if frame is None:
                break
            now_path = os.getcwd()
            result = frame



            videoWriter.write(result)
            key = cv2.waitKey(1) & 0xFF
            if key == ord("q"):
                break
    return s_path

